#!/usr/bin/env python

import math
import Image
import serial
import struct
import sys
import time

class Robot:
    """Robot class. For each instruction you send to the robot, you
    need to empty its output buffer using parse_output().  
    Note that using serial.readlines() has a timeout, making the
    robot seemingly slow to respond. Use serial.readline() instead."""
    def __init__(self):
        self.axle_diameter = 53.0
        self.axle_circum = math.pi * self.axle_diameter
        self.wheel_diameter = 41.0
        self.wheel_circum = math.pi * self.wheel_diameter
        self.timesteps = self.wheel_circum / 1000

    def com_handler(method):
        """Handle attempts to write() after being disconnected."""
        def safe(*args, **kwargs):
            try:
                method(*args, **kwargs)
            except OSError:
                sys.stderr.write("Lost connection!\n")
                sys.exit(1)
        return safe

    def connect(self, port, bytesize = 8, baudrate = 115200, timeout = 1):
        """Connect to the robot through a virtual serial-port (Bluetooth)."""
        try:
            self.com = serial.Serial(port, baudrate,
                                     bytesize, timeout = timeout)
        except serial.serialutil.SerialException:
            sys.stderr.write("*** No connection. Robot appears to be offline\n")
            sys.exit(1)
        print "*** Connected to robot"
        self.wake_up()

    def disconnect(self):
        """Disconnect from robot."""
        print "*** Disconnected"
        self.com.close()

    def parse_output(self, data):
        """Parse the output from the robot."""
        pass

    @com_handler
    def wake_up(self):
        """The robot never executes the first instruction it receives."""
        self.com.write("\n")
        output = self.parse_output(self.com.readline())

    def set_motor_speeds(self, left, right):
        """Motor speed instruction."""
        return "D,%d,%d\n" % (int(left), int(right))

    @com_handler
    def move(self, direction, duration):
        """Move in the specified direction for n seconds."""
        if direction == "forward":
            instruction = self.set_motor_speeds(500, 500)
        elif direction == "backward":
            instruction = self.set_motor_speeds(-500, -500)
        else:
            return            
        self.com.write(instruction)
        output = self.parse_output(self.com.readline())
        time.sleep(duration)
        self.stop()

    @com_handler
    def spin(self, angle, direction):
        """Spin clockwise or counter-clockwise."""
        angle = angle / 360.0
        distance = self.axle_circum * angle        
        spin_time = distance / self.timesteps / 1000.0
        if direction == "left":
            instruction = self.set_motor_speeds(-1000, 1000)
        elif direction == "right":
            instruction = self.set_motor_speeds(1000, -1000)            
        self.com.write(instruction)
        time.sleep(spin_time)
        self.stop()

    def turn_left(self):
        """Turn 90 degrees left."""
        self.spin(90, "left")

    def turn_right(self):
        """Turn 90 degrees right."""
        self.spin(90, "right")

    @com_handler
    def stop(self):
        """Set motor speed to zero."""
        instruction = self.set_motor_speeds(0, 0)
        self.com.write(instruction)
        output = self.parse_output(self.com.readline())

    def capture_image(self):
        """Capture raw 5:6:5 RGB image data from the camera
        and convert to PNG for compatbility."""
        try:
            instruction = struct.pack(">bb", - ord("I"), 0)
            self.com.write(instruction)
            output = self.com.read(999999)
        except OSError: 
            sys.stderr.write("Lost connection!\n")
            sys.exit(1)
        try:
            buffer = Image.frombuffer("RGB", (40, 40), output,
                                      "raw", "BGR;16", 0, 1)
        except ValueError, msg:
            print "Bad connection: %s" % msg
            return
        buffer = buffer.rotate(180)
        buffer = buffer.resize((300, 300))
        buffer.show()
        return buffer

def parse_input(robot):
    """Read user commands and redirect to robot."""
    while True:
        input = raw_input()
        if "quit" in input:
            return
        elif "forward" in input:
            try:
                cmd, duration = input.split()
            except ValueError:
                print "Syntax: forward <seconds>"
            robot.move("forward", float(duration))
        elif "backward" in input:
            try:
                cmd, duration = input.split()
            except ValueError:
                print "Syntax: backward <seconds>"
            robot.move("backward", float(duration))
        elif "right" in input:
            robot.turn_right()
        elif "left" in input:
            robot.turn_left()
        elif "snapshot" in input:
            robot.capture_image()
        else:
            print "Unknown command. Type \"quit\" to exit."

if __name__ == "__main__":
    r = Robot()
    r.connect("/dev/tty.e-puck_1251-COM1-1")
    parse_input(r)
    r.disconnect()
